
#include "tensor.h"
#include "math_ops.h"
#include <iostream>
#include <hip/hip_runtime.h>

int main() {
    // 在CPU上分配两个张量用于测试，这里简单示例都是一维张量，元素个数为10，数据类型为FLOAT32
    Tensor tensor_a = allocateTensorOnCPU({10}, DataType::FLOAT32);
    Tensor tensor_b = allocateTensorOnCPU({10}, DataType::FLOAT32);

    // 简单初始化一下数据，这里示例使用float* 来访问数据进行赋值（仅用于测试示例）
    float* data_a = static_cast<float*>(tensor_a.data);
    float* data_b = static_cast<float*>(tensor_b.data);
    for (int i = 0; i < 10; ++i) {
        data_a[i] = i * 1.1;
        data_b[i] = 2.3 * i;
    }

    // 将数据从CPU拷贝到GPU上，分别为两个输入张量分配GPU内存并拷贝数据
    Tensor tensor_a_gpu = allocateTensorOnGPU({10}, DataType::FLOAT32);
    Tensor tensor_b_gpu = allocateTensorOnGPU({10}, DataType::FLOAT32);
    copyToGPU(tensor_a, tensor_a_gpu);
    copyToGPU(tensor_b, tensor_b_gpu);

    // 调用加法运算函数进行GPU上的加法操作
    Tensor result_gpu = add_gpu(tensor_a_gpu, tensor_b_gpu);
    
    // 在CPU上分配一个张量用于接收GPU计算的结果
    Tensor result_cpu = allocateTensorOnCPU({10}, DataType::FLOAT32);

    // 将GPU计算结果拷贝回CPU
    copyToGPU(result_gpu, result_cpu);

    // 输出结果验证（简单打印结果，实际应用中可能有更复杂的验证方式）
    float* result_data = static_cast<float*>(result_cpu.data);
    std::cout << "加法运算结果: ";
    for (int i = 0; i < 10; ++i) {
        std::cout << result_data[i] << " ";
    }
    std::cout << std::endl;

    // 释放内存，按照分配顺序的逆序释放各个张量占用的内存
    freeTensor(result_cpu);
    freeTensor(result_gpu);
    freeTensor(tensor_b_gpu);
    freeTensor(tensor_a_gpu);
    freeTensor(tensor_b);
    freeTensor(tensor_a);
    
    return 0;
}

